/* * * * *
 *  predict.cpp
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  Copyright (C) 2012, Rongkai Xia, shinekaixia@gmail.com
 *
 * * * * */

#include <iostream>
#include "RandomForestRegressor.h"
#include "RandomForestClassifier.h"
#include "GBMRegressor.h"
#include "GBMClassifier.h"
#include "ReadData.h"
#include "cmdline.h"

static int support_alg_len=4;
static string support_alg[4]={
    "GBMRegressor",
    "GBMClassifier",
    "RandomForestRegressor",
    "RandomForestClassifier"
};

static int predict_param_len=3;
static string predict_param[3]={
    "test_x_fn",
    "model_fn",
    "prediction_fn"
};

static string predict_help[3]={
    "path to the data point file of training data.",
    "path to the model file.",
    "path to the predcition file to store the prediction."
};

static void print_help(){
    cmdline cmd;
    printf("Usage:\n");
    cmd.register_help(predict_param, predict_help, predict_param_len);
    cmd.print_help();
}

bool do_predict(cmdline& cmd){
    string test_x_fn,model_fn,prediction_fn,model_str,param;
    int i;
    Data test;
    REAL* pred;
    if ( !cmd.check_param(predict_param, predict_param_len)){
        return false;
    }
    if (!(cmd.has_param("test_x_fn") && cmd.has_param("model_fn") && cmd.has_param("prediction_fn")) ){
        fprintf(stderr, "Error: train_x_fn,model_fn,and prediction_fn must be set.\n");
        fprintf(stderr, "Please type -help for help.\n");  
        return false;
    }
    test_x_fn=cmd.get_value("test_x_fn");
    model_fn=cmd.get_value("model_fn");
    prediction_fn=cmd.get_value("prediction_fn");
    //test path correctness,try create prediction file 
    FILE* pred_fp=fopen(prediction_fn.c_str(), "w");
    if (!pred_fp) {
        fprintf(stderr, "Error: try create prediction file %s error.Please check your prediction file path.\n",prediction_fn.c_str());
        fclose(pred_fp);
        return false;
    }
    fclose(pred_fp);
    fstream model_input(model_fn.c_str());
    if (model_input.fail()) {
        fprintf(stderr, "Error: cannot open model file %s.\n",model_fn.c_str());
        return false;
    }
    model_input>>model_str;
    //test model correctness
    for (i=0; i<support_alg_len; i++) {
        if (model_str==support_alg[i]) {
            break;
        }
    }
    if (i==support_alg_len) {
        fprintf(stderr, "Error: Unknow model file type %s,the model file is corrupt.\n",model_str.c_str());
        return false;
    }
    model_input.close();
    
    //read test files.
    if (!readData(test, test_x_fn.c_str())){
        return false;
    }
    if (test.n_samples==0) {
        fprintf(stderr, "Error: number of samples in test file is 0.\n");
        return false;
    }
    pred=new REAL[test.n_samples];
    
    //read model file
    if (model_str=="GBMRegressor") {
        GBMRegressor m;
        if (!m.load_model(model_fn.c_str())){
            return false;
        }
        if (test.n_features!=m.n_features) {
            fprintf(stderr, "Error: number of features in test file and train file are different,train file has %d features,but test file has %d features.Please check your input files\n",m.n_features,test.n_features);
            return false;
        }
        m.predict(test.X, pred, test.n_samples, test.n_features);
    }else if (model_str=="GBMClassifier"){
        GBMClassifier m;
        if (!m.load_model(model_fn.c_str())){
            return false;
        }
        if (test.n_features!=m.n_features) {
            fprintf(stderr, "Error: number of features in test file and train file are different,train file has %d features,but test file has %d features.Please check your input files\n",m.n_features,test.n_features);
            return false;
        }
        m.predict(test.X, pred, test.n_samples, test.n_features);
    }else if (model_str=="RandomForestRegressor"){
        RandomForestRegressor m;
        if (!m.load_model(model_fn.c_str())){
            return false;
        }
        if (test.n_features!=m.n_features) {
            fprintf(stderr, "Error: number of features in test file and train file are different,train file has %d features,but test file has %d features.Please check your input files\n",m.n_features,test.n_features);
            return false;
        }
        m.predict(test.X, pred, test.n_samples, test.n_features);
    }else if (model_str=="RandomForestClassifier"){
        RandomForestClassifier m;
        if (!m.load_model(model_fn.c_str())){
            return false;
        }
        if (test.n_features!=m.n_features) {
            fprintf(stderr, "Error: number of features in test file and train file are different,train file has %d features,but test file has %d features.Please check your input files\n",m.n_features,test.n_features);
            return false;
        }
        m.predict(test.X, pred, test.n_samples, test.n_features);
    }else {
        fprintf(stderr, "Error: Unknow model file type %s,the model file is corrupt.\n",model_str.c_str());
        return false;
    }
    pred_fp=fopen(prediction_fn.c_str(), "w");
    for (int i=0; i<test.n_samples; i++) {
//        fprintf(stderr, "%.10lg\n",pred[i]);
        fprintf(pred_fp, "%.10lg\n",pred[i]);
    }
    fclose(pred_fp);
    return true;
}
int main(int argc, const char * argv[])
{
    srand((uint)time(NULL));
    seedMT(2*rand()+1);
    cmdline cmd;
    if (!cmd.parse_cmdline(argc, argv)){
        fprintf(stderr, "Please type -help for help.\n");
        return 1;
    }
    if (cmd.has_param("help")) {
        print_help();
        return 1;
    }
    if (!do_predict(cmd)){
        return 1;
    }
    return 0;
}

